#include <inc/assert.h>
#include <inc/log.h>
#include <kern/env.h>
#include <kern/pmap.h>
#include <kern/monitor.h>

// Calculate next environment to pick up by round-robin.
// Will only pick envs[1]~envs[NENV-1]
static inline struct Env *
rr_next(struct Env *env)
{
	// If env is invalid, then we just return envs
	// otherwise envs + (env - envs + 1) % NENV
	if (env < envs || env >= envs + NENV)
		return envs;
	else
		return envs + (env - envs + 1) % NENV;
}

// Choose a user environment to run and run it.
void
rr_sched_yield(void)
{
	// Implement simple round-robin scheduling.
	// Search through 'envs' for a runnable environment,
	// in circular fashion starting after the previously running env,
	// and switch to the first such environment found.
	// It's OK to choose the previously running env if no other env
	// is runnable.
	// But never choose envs[0], the idle environment,
	// unless NOTHING else is runnable.

	// LAB 4: Your code here.

	// Since all user environment switches are done by this
	// function, env_to_run keeps the last run environment.
	static struct Env *env_to_run = NULL;
	struct Env *end_env;

	if (env_to_run) {
		end_env = env_to_run;
		env_to_run = rr_next(env_to_run);
	} else if (curenv) {
		env_to_run = rr_next(curenv);
		end_env = curenv;
	} else {
		env_to_run = envs + 1;
		end_env = envs;
	}
	while (env_to_run != end_env) {
		if (env_to_run->env_status == ENV_RUNNABLE && env_to_run != envs) {
//			cprintf("yield %x %x from %x\n", env_to_run->env_id, env_to_run->env_tf.tf_eip, end_env->env_tf.tf_eip);
			env_run(env_to_run);
		}
		env_to_run = rr_next(env_to_run);
	}
	if (env_to_run->env_status == ENV_RUNNABLE) {
//		cprintf("yield %x %x from %x\n", env_to_run->env_id, env_to_run->env_tf.tf_eip, end_env->env_tf.tf_eip);
		env_run(env_to_run);
	}

	// Run the special idle environment when nothing else is runnable.
	if (envs[0].env_status == ENV_RUNNABLE)
		env_run(&envs[0]);
	else {
		cprintf("Destroyed all environments - nothing more to do!\n");
		while (1)
			monitor(NULL);
	}
}

// Lottery sched
// FIXME: slow, non-working with big number of envs
//Park-Miller pseudo random number generator
static inline uint32_t
lcg_rand(void)
{
	static uint32_t a = 0xdeadbeef;
        a = ((unsigned long)a * 279470273UL) % 4294967291UL;
	return a;
}

void
lottery_sched_yield(void)
{
	// no one has a ticket, this shouldn't happen
	uint32_t total_tickets = 0;
	uint32_t winner;
	struct Env *e, *ef = NULL;
	// collect the tickets
	for (e = envs + 1; e < envs + NENV; ++e)
		if (e->env_status == ENV_RUNNABLE) {
			ef = ef ? ef : e;
			total_tickets += e->env_ticket;
		}
	// If no runnable env, then run the special idle environment when nothing else is runnable.
	if (total_tickets == 0) {
		if (envs[0].env_status == ENV_RUNNABLE)
			env_run(&envs[0]);
		else {
			cprintf("Destroyed all environments - nothing more to do!\n");
			while (1)
				monitor(NULL);
		}
	}
	// otherwise choose a winner
	winner = lcg_rand() % total_tickets;
	// start from the first runnable env

	/*
	cprintf("T %d W %d\n", total_tickets, winner);
	for (e = ef; e < envs + NENV; ++e)
		if (e->env_status == ENV_RUNNABLE)
			cprintf("%x %d\n", e->env_id, e->env_ticket);
	*/
	e = ef;
	while (winner) {
		if (e->env_status == ENV_RUNNABLE) {
			if (winner >= e->env_ticket)
				winner -= e->env_ticket;
			else {
				winner = 0;
				break;
			}
		}
		do
			++e;
		while (e->env_status != ENV_RUNNABLE);
	}

//	cprintf("Will run %x\n", e->env_id);

	env_run(e);
}

void
sched_yield(void)
{
//	lottery_sched_yield();
	rr_sched_yield();
}
